package port

import (
	"net"
	"strconv"
	"time"

	"github.com/miekg/dns"
)

// CheckUDP 检测UDP端口，返回检测结果
func (pc *PortChecker) CheckUDP(host string, port int) *PortCheckResult {
	result := &PortCheckResult{
		Host:   host,
		Port:   port,
		Sent:   pc.Count,
		Delays: make([]time.Duration, 0, pc.Count),
	}
	address := net.JoinHostPort(host, strconv.Itoa(port))
	var totalBytes int64
	for i := 0; i < pc.Count; i++ {
		var success bool
		var cost time.Duration
		if port == 53 {
			// 发送标准DNS查询包
			m := new(dns.Msg)
			m.SetQuestion("baidu.com.", dns.TypeA)
			c := new(dns.Client)
			c.Timeout = time.Duration(pc.Timeout) * time.Second
			_, rtt, err := c.Exchange(m, address)
			cost = rtt
			if err == nil {
				success = true
				totalBytes += int64(len(m.Question))
			}
		} else {
			conn, err := net.DialTimeout("udp", address, time.Duration(pc.Timeout)*time.Second)
			if err != nil {
				result.Failed++
				continue
			}
			start := time.Now()
			conn.SetDeadline(time.Now().Add(time.Duration(pc.Timeout) * time.Second))
			msg := []byte("gns_port_test")
			_, err = conn.Write(msg)
			if err != nil {
				conn.Close()
				result.Failed++
				continue
			}
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			cost = time.Since(start)
			if err == nil && n > 0 {
				success = true
				totalBytes += int64(n)
			}
			conn.Close()
		}
		if success {
			result.Success++
			result.Delays = append(result.Delays, cost)
		} else {
			result.Failed++
		}
		time.Sleep(200 * time.Millisecond)
	}
	if result.Success > 0 && len(result.Delays) > 0 {
		result.CalculateStats()
		if result.AvgDelay > 0 {
			result.Speed = float64(totalBytes) / result.AvgDelay.Seconds()
		}
	}
	return result
}
